# -*- coding:utf-8  -*-
# @Author: ChenYangMing
# @Time: 2024-09-25

from datetime import timedelta
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.preprocessing import MinMaxScaler
import torch
from torch.utils.data import DataLoader, TensorDataset

from Common.file_utils import FileUtils
from Common.data_utils import DataUtils
from config import TaskConfig


class TaskDatasets(DataUtils, FileUtils):
    """
    任务数据集类，用于创建和处理城市网络流量长时预测任务所需的数据集。
    """
    def __init__(
            self,
            logger,
    ):
        """
        初始化任务数据集类。
        :param logger: 日志记录器对象。
        """
        self.logger = logger
        self.sequence_length = TaskConfig.HISTORY_LENGTH
        self.prediction_length = TaskConfig.PREDICTION_LENGTH
        self.datasets = {}
        self.input_dim, self.output_dim = 0, 0
        self.scaler = MinMaxScaler()

    def create_dataset(
            self,
            data_path: str,
            holidays_data_path: str,
    ):
        """
        构造训练集，包括训练数据和标签
        :param data_path: str, 流量数据存储路径。
        :param holidays_data_path: str, 节假日数据存储路径。
        """
        traffic_df = pd.read_csv(data_path)[['时间', '地市', '流量']]
        traffic_df = traffic_df.rename(columns={'时间': 'date', '地市': 'city', '流量': 'traffic'})
        traffic_df = self.add_time_fea(traffic_df)

        holidays_df = pd.read_csv(holidays_data_path)
        traffic_df = self.add_holiday_fea(holidays_df, traffic_df)
        traffic_df = self.fill_value(traffic_df)
        # self.analysis_fea(traffic_df)
        traffic_df['traffic'] = self.scaler.fit_transform(traffic_df['traffic'].values.reshape(-1, 1))
        for city in ['A', 'B', 'C']:
            city_df = traffic_df[traffic_df['city'] == city][['year', 'month', 'hour', 'traffic']]
            city_df.reset_index(drop=True, inplace=True)
            X_train, y_train = self.create_time_series(city_df, self.sequence_length, self.prediction_length)
            # X: (样本数, 时间步数, 特征数), y: (样本数, 预测时间步数)
            self.logger.info(f"city {city}: shape of input - {X_train.shape}, shape of output - {y_train.shape}")
            self.input_dim = X_train.shape[2]
            self.output_dim = y_train.shape[1]
            X_train_tensor = torch.tensor(X_train, dtype=torch.float32)
            y_train_tensor = torch.tensor(y_train, dtype=torch.float32)
            dataset = TensorDataset(X_train_tensor, y_train_tensor)
            dataloader = DataLoader(dataset, batch_size=TaskConfig.BATCH_SIZE, shuffle=True)
            X_test = self.prepare_test_data(city_df, self.sequence_length)
            X_test_tensor = torch.tensor(X_test, dtype=torch.float32)
            self.datasets[city] = {'train': dataloader, 'test': X_test_tensor}

    def add_time_fea(
            self,
            df: pd.DataFrame,
    ):
        """
        添加时间特征：年、月、日、时、季度、星期几。
        :param df: pd.DataFrame, 流量数据。
        :return: pd.DataFrame, 添加时间特征后的流量数据。
        """
        df['date'] = pd.to_datetime(df['date'])
        df['year'] = df['date'].dt.year
        df['month'] = df['date'].dt.month
        df['day'] = df['date'].dt.day
        df['hour'] = df['date'].dt.hour
        df['quarter'] = df['date'].dt.quarter  # 添加季度特征
        df['weekday'] = df['date'].dt.weekday  # 0 表示星期一，1 表示星期二，以此类推，6 表示星期日
        self.logger.info('添加时间特征：年、月、日、时、季度、星期几')
        return df

    def add_holiday_fea(
            self,
            holidays_df,
            df,
    ):
        """
        添加是否节日特征。
        :param holidays_df: pd.DataFrame, 节假日数据。
        :param df: pd.DataFrame, 流量数据。
        :return: pd.DataFrame, 添加是否节日特征后的流量数据。
        """
        holidays_timestamps = []
        for i in range(holidays_df.shape[0]):
            for j in range(1, holidays_df.shape[1], 2):
                start_time = pd.to_datetime(holidays_df.iloc[i, j])
                end_time = pd.to_datetime(holidays_df.iloc[i, j + 1])
                timestamps = [start_time + timedelta(hours=i) for i in
                              range(int((end_time - start_time).total_seconds() / 3600) + 1)]
                holidays_timestamps += timestamps
        df['is_holiday'] = df['date'].isin(holidays_timestamps)
        self.logger.info('添加是否节日特征')
        return df

    def fill_value(
            self,
            df,
    ):
        """
        若流量数据中存在缺失值，则用同 city、year、month、weekday、is_holiday、hour 的流量均值进行填充。
        :param df: pd.DataFrame, 流量数据。
        :return: pd.DataFrame, 填充缺失值后的流量数据。
        """
        num_null = df['traffic'].isnull().sum()
        num_zero = df['traffic'].eq(0).sum()
        if num_null + num_zero > 0:
            grouped_means = df.groupby(
                ['city', 'year', 'month', 'weekday', 'is_holiday', 'hour'])['traffic'].mean().reset_index()
            grouped_means.rename(columns={'traffic': 'mean_traffic'}, inplace=True)
            df = pd.merge(df, grouped_means, on=['city', 'year', 'month', 'weekday', 'is_holiday', 'hour'], how='left')
            df['traffic'] = df['traffic'].replace(0, np.nan)
            df['traffic'].fillna(df['mean_traffic'], inplace=True)
            df.drop(columns=['mean_traffic'], inplace=True)
            self.logger.info('填充缺失值')
        return df

    @staticmethod
    def analysis_fea(
            df,
    ):
        """
        通过箱型图分析不同特征对流量的影响。
        :param df: pd.DataFrame, 流量数据。
        """
        fea_list = ['year', 'month', 'day', 'hour', 'weekday', 'quarter', 'is_holiday', 'city']
        for fea in fea_list:
            plt.figure(figsize=(15, 10))
            sns.boxplot(x=fea, y='traffic', data=df)
            plt.title(f'Traffic Distribution by {fea.capitalize()}')
            plt.xlabel(fea.capitalize())
            plt.ylabel('Traffic Volume')
            plt.xticks(rotation=45)
            plt.tight_layout()  # 调整布局，避免图像被截断
            plt.show()
            plt.close()

    @staticmethod
    def create_time_series(
            x,
            hl,
            pl,
    ):
        """
        将 x 转换为时间序列数据。
        :param x: pd.DataFrame, 输入特征数据。
        :param hl: int, 历史窗口长度。
        :param pl: int, 预测长度。
        :return: tuple, 包含时间序列的元组 (X, y)，其中 X 是特征集，y 是目标集。
        """
        X, y = [], []
        total_time_window = hl + pl
        loop_count = len(x) - total_time_window  # 循环次数
        for i in range(loop_count):
            X.append(x.iloc[i:(i + hl)].values)  # 历史数据
            y.append(x.iloc[i + hl:i + total_time_window]['traffic'].values)  # 目标
        return np.array(X), np.array(y)

    @staticmethod
    def prepare_test_data(
            df,
            hw,
    ):
        """
        准备测试数据。
        :param df: pd.DataFrame, 输入特征数据。
        :param hw: int, 历史窗口长度。
        :return: np.ndarray, 准备的测试数据。
        """
        x = [df.iloc[len(df) - hw: len(df)].values]
        return np.array(x)


# 调用示例
# file_path = 'data.xlsx'
# sheet_name = 'Sheet1'
# processor = TaskDatasets()  # 创建 TaskDatasets 实例
# data = processor.read_file(file_path, file_type='xlsx', sheet_name=sheet_name)
# processed_data = processor.process_data(data)  # 处理数据
# output_file_path = 'processed_data.csv'
# processor.write_file(output_file_path, processed_data, file_type='csv')
